import torch
import torch.nn as nn
from transformers import BertModel
import torch.nn.functional as F

class BertSim(nn.Module):

    def __init__(self, bert_config, hidden_size, n_classes, dropout):
        super(BertSim, self).__init__()
        self.hidden_size = hidden_size
        self.bert = BertModel.from_pretrained(bert_config)
        self.linear = nn.Linear(self.hidden_size, n_classes)
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, sentence, attention_mask=None, token_type_ids=None):
        embeds = self.bert(sentence, attention_mask=attention_mask, token_type_ids=token_type_ids)
        embeds_cls = embeds[0][:, 0, :] # CLS的向量输出
        embeds_cls_dropout = self.dropout(embeds_cls)
        out = self.linear(embeds_cls_dropout)
        return out # [batch_size, num_classes]
